// Copyright (c) 2020 Presto Labs Pte. Ltd.
// Author: xguo

#pragma once

#include <functional>
#include <memory>
#include <queue>
#include <string>
#include <vector>
#include <boost/asio.hpp>
#include <boost/asio/spawn.hpp>

#include "coin2/base/log.h"
#include "coin2/base/time.h"

struct Task {
  using TaskFunc = std::function<void()>;

  Task(TaskFunc task_func, int64_t period, int64_t next_run_time, const std::string& task_name)
      : task_func(task_func), period(period), next_run_time(next_run_time), task_name(task_name) {}

  void ScheduleNextRun() {
    const auto curr_ts = GetCurrentTimestamp();
    next_run_time = curr_ts + period;
  }

  void ScheduleNextRun(const int64_t curr_ts) {
    CHECK(period > 0) << "Invalid period: " << period;
    while (next_run_time <= curr_ts) {
      next_run_time += period;
    }
  }

  void Run() const { task_func(); }

  TaskFunc task_func;
  int64_t period;
  int64_t next_run_time;
  std::string task_name;
};

struct TaskCompare {
  bool operator()(const Task* lhs, const Task* rhs) {
    return lhs->next_run_time > rhs->next_run_time;  // Small one has higher priority
  }
};

class Scheduler : public std::enable_shared_from_this<Scheduler> {
 public:
  using TaskFunc = Task::TaskFunc;

 public:
  explicit Scheduler(boost::asio::io_context*);
  explicit Scheduler(boost::asio::io_context::strand*);

  void AddPeriodicTask(TaskFunc, int64_t period, const std::string& task_name);
  void AddPeriodicTask(TaskFunc, int64_t period, const std::string& task_name, int64_t delay);
  void CallLater(TaskFunc, int64_t delay, const std::string& task_name);
  void StartImpl(boost::asio::yield_context yield_context);
  void Start();
  void Stop();
  void SetCheckPeriod(int64_t check_period_ns);

 private:
  std::priority_queue<Task*, std::vector<Task*>, TaskCompare> task_queue_;
  std::vector<std::unique_ptr<Task>> tasks_;
  std::vector<Task*> task_vec_;
  boost::asio::io_context* ioc_;
  boost::asio::io_context::strand* strand_;
  bool running_{false};
  int64_t check_period_ns_{10'000'000};  // 10ms
};
